from yk_tools import some_def
from yk_tools import halo_decorator 
some_def.os.chdir("/Pub/Users/wangyk/project/Poroject/F241121002_20250210/")
import scanpy as sc
import os
import tqdm

# Define the directories containing your 10x data for each sample
# Replace with the actual paths to your directories

data_directories = ["data/GSE175453/" + i for i in os.listdir("data/GSE175453")]
# --- Step 1: Read each 10x dataset and store them in a list ---
adatas = []
sample_names = os.listdir("data/GSE175453") # To store sample names for annotation

for directory in tqdm.tqdm(data_directories):
    adata = sc.read_10x_mtx(directory,gex_only = True)
    sc.pp.filter_cells(adata,min_genes=200)
    sc.pp.filter_genes(adata,min_cells=2)
    adata.var_names_make_unique()  # Ensure variable names are unique
    adata.obs['sample'] = os.path.basename(directory)  # Add 'sample' column to .obs dataframe
    adata.obs.index = adata.obs.index + os.path.basename(directory)
    adatas.append(adata)

# --- Step 2: Concatenate AnnData objects into a single AnnData object ---
# Use scanpy's concatenate function to merge the list of AnnData objects
adata_merged = sc.concat(
    adatas,
    label="sample",  # Column name to store batch/sample information in .obs
    keys=sample_names, # Values to use for the 'sample' column, corresponding to each input AnnData object
    merge="same"      # Ensure variables (genes) are the same across datasets. If genes differ, use 'unique' or other options.
)

# --- Step 3: (Optional) Print basic information about the merged AnnData object ---
print(adata_merged)
print(adata_merged.obs.head()) # Show the first few rows of the .obs dataframe, including the 'sample' column

# --- Step 4: (Optional) Further processing and analysis in Scanpy ---
# You can now proceed with your single-cell analysis pipeline using 'adata_merged'
# For example, basic preprocessing steps:

adata_merged.var['mt'] = adata_merged.var_names.str.startswith('MT-')  # annotate mitochondrial genes as 'mt'
sc.pp.calculate_qc_metrics(adata_merged, qc_vars=['mt'], percent_top=None, log1p=False, inplace=True)

adata_merged = adata_merged[adata_merged.obs.pct_counts_mt < 20, :] # Filter cells based on mitochondrial counts
adata_merged = adata_merged[adata_merged.obs.n_genes_by_counts > 200, :] # Filter cells based on gene counts

sc.pp.normalize_total(adata_merged, target_sum=1e4)
sc.pp.log1p(adata_merged)
sc.pp.highly_variable_genes(adata_merged, min_mean=0.0125, max_mean=3, min_disp=0.5)
adata_merged.raw = adata_merged # Store raw counts before filtering highly variable genes
adata_merged = adata_merged[:, adata_merged.var.highly_variable] # Filter to highly variable genes

sc.pp.scale(adata_merged, max_value=10)
sc.tl.pca(adata_merged, svd_solver='arpack')
sc.pp.neighbors(adata_merged, n_neighbors=10, n_pcs=40)
sc.tl.umap(adata_merged)
sc.tl.leiden(adata_merged)
sc.pl.umap(adata_merged, color=['leiden', 'sample']) # Visualize UMAP colored by cluster and sample


# --- Step 5: (Optional) Save the merged AnnData object for later use ---
# adata_merged.write('./merged_anndata.h5ad')